package com.nanohadoop.hdfs.utils;

import com.nanohadoop.hdfs.NameNode.DataNodeInfo;
import com.nanohadoop.hdfs.NameNode.BlockLocation;
import com.nanohadoop.hdfs.HDFSClient.FileStatus;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.*;
import java.net.*;
import java.util.*;

/**
 * 网络通信工具类
 */
public class NetworkUtils {
    private static final Logger log = LoggerFactory.getLogger(NetworkUtils.class);
    private static final int SOCKET_TIMEOUT = 5000;

    public boolean isReachable(String host, int port) {
        try (Socket socket = new Socket()) {
            socket.connect(new InetSocketAddress(host, port), SOCKET_TIMEOUT);
            return true;
        } catch (IOException e) {
            return false;
        }
    }

    public boolean requestNameNode(String host, int port, String command, String path) 
            throws IOException {
        try (Socket socket = new Socket(host, port);
             DataOutputStream out = new DataOutputStream(socket.getOutputStream());
             DataInputStream in = new DataInputStream(socket.getInputStream())) {
            
            // 发送请求
            out.writeUTF(command);
            out.writeUTF(path);
            out.flush();
            
            // 读取响应
            return in.readBoolean();
        }
    }

    public List<DataNodeInfo> requestBlockLocations(String host, int port, 
            String filename, long blockId, long blockLength) throws IOException {
        try (Socket socket = new Socket(host, port);
             DataOutputStream out = new DataOutputStream(socket.getOutputStream());
             DataInputStream in = new DataInputStream(socket.getInputStream())) {
            
            out.writeUTF("GET_BLOCK_LOCATIONS");
            out.writeUTF(filename);
            out.writeLong(blockId);
            out.writeLong(blockLength);
            out.flush();
            
            // 读取DataNode列表
            int count = in.readInt();
            List<DataNodeInfo> locations = new ArrayList<>();
            for (int i = 0; i < count; i++) {
                String id = in.readUTF();
                String dataNodeHost = in.readUTF();
                int dataNodePort = in.readInt();
                locations.add(new DataNodeInfo(id, dataNodeHost, dataNodePort));
            }
            return locations;
        }
    }

    public List<BlockLocation> getFileBlockLocations(String host, int port, String path) 
            throws IOException {
        try (Socket socket = new Socket(host, port);
             DataOutputStream out = new DataOutputStream(socket.getOutputStream());
             DataInputStream in = new DataInputStream(socket.getInputStream())) {
            
            log.debug("Requesting blocks for file: " + path);
            out.writeUTF("GET_FILE_BLOCKS");
            out.writeUTF(path);
            out.flush();
            
            // 读取块位置信息
            int count = in.readInt();
            List<BlockLocation> locations = new ArrayList<>();
            for (int i = 0; i < count; i++) {
                long blockId = in.readLong();
                int nodeCount = in.readInt();
                List<DataNodeInfo> nodes = new ArrayList<>();
                for (int j = 0; j < nodeCount; j++) {
                    String nodeId = in.readUTF();
                    String dataNodeHost = in.readUTF();
                    int dataNodePort = in.readInt();
                    nodes.add(new DataNodeInfo(nodeId, dataNodeHost, dataNodePort));
                }
                locations.add(new BlockLocation(blockId, nodes));
            }
            return locations;
        }
    }

    public void writeBlock(String datanodeId, long blockId, byte[] data) throws IOException {
        if (datanodeId == null || !datanodeId.contains(":")) {
            throw new IOException("Invalid DataNode ID format: " + datanodeId + ". Expected format: host:port");
        }
        
        String[] parts = datanodeId.split(":");
        if (parts.length != 2) {
            throw new IOException("Invalid DataNode ID format: " + datanodeId + ". Expected format: host:port");
        }
        
        String host = parts[0];
        int port;
        try {
            port = Integer.parseInt(parts[1]);
        } catch (NumberFormatException e) {
            throw new IOException("Invalid port in DataNode ID: " + datanodeId, e);
        }
        
        log.debug("Writing block " + blockId + " to DataNode: " + datanodeId);
        try (Socket socket = new Socket(host, port);
             DataOutputStream out = new DataOutputStream(socket.getOutputStream())) {
            
            out.writeUTF("WRITE");
            out.writeLong(blockId);
            out.writeInt(data.length);
            out.write(data);
            out.flush();
            
            // 等待写入确认
            try (DataInputStream in = new DataInputStream(socket.getInputStream())) {
                boolean success = in.readBoolean();
                if (!success) {
                    throw new IOException("Failed to write block " + blockId + " to DataNode " + datanodeId);
                }
            }
        }
        log.debug("Successfully wrote block " + blockId + " to DataNode: " + datanodeId);
    }

    public byte[] readBlock(String datanodeId, long blockId) throws IOException {
        if (datanodeId == null || !datanodeId.contains(":")) {
            throw new IOException("Invalid DataNode ID format: " + datanodeId + ". Expected format: host:port");
        }
        
        String[] parts = datanodeId.split(":");
        String host = parts[0];
        int port;
        try {
            port = Integer.parseInt(parts[1]);
        } catch (NumberFormatException e) {
            throw new IOException("Invalid port in DataNode ID: " + datanodeId, e);
        }
        
        log.debug("Connecting to DataNode " + datanodeId + " to read block " + blockId);
        try (Socket socket = new Socket(host, port);
             DataOutputStream out = new DataOutputStream(socket.getOutputStream());
             DataInputStream in = new DataInputStream(socket.getInputStream())) {
            
            out.writeUTF("READ");
            out.writeLong(blockId);
            out.flush();
            
            int length = in.readInt();
            log.debug("DataNode " + datanodeId + " reported block " + blockId + 
                " length as " + length + " bytes");
                
            if (length > 0) {
                byte[] data = new byte[length];
                in.readFully(data);
                log.debug("Successfully read " + data.length + " bytes from block " + 
                    blockId + " on DataNode " + datanodeId);
                return data;
            }
            log.debug("Warning: DataNode " + datanodeId + " returned empty block " + blockId);
            return new byte[0];
        } catch (IOException e) {
            log.debug("Failed to read block " + blockId + " from DataNode " + datanodeId + 
                ": " + e.getMessage());
            throw e;
        }
    }

    public byte[] fetchBlock(String sourceDataNode, long blockId) throws IOException {
        return readBlock(sourceDataNode, blockId);
    }

    public FileStatus getFileStatus(String host, int port, String path) throws IOException {
        try (Socket socket = new Socket(host, port);
             DataOutputStream out = new DataOutputStream(socket.getOutputStream());
             DataInputStream in = new DataInputStream(socket.getInputStream())) {
            
            out.writeUTF("GET_FILE_STATUS");
            out.writeUTF(path);
            out.flush();
            
            boolean exists = in.readBoolean();
            if (!exists) {
                throw new FileNotFoundException(path);
            }
            
            return new FileStatus(
                path,
                in.readLong(),     // length
                in.readBoolean(),  // isDirectory
                in.readInt(),      // replication
                in.readLong()      // blockSize
            );
        }
    }

    public List<FileStatus> listFiles(String host, int port, String path) throws IOException {
        try (Socket socket = new Socket(host, port);
             DataOutputStream out = new DataOutputStream(socket.getOutputStream());
             DataInputStream in = new DataInputStream(socket.getInputStream())) {
            
            out.writeUTF("LIST_FILES");
            out.writeUTF(path);
            out.flush();
            
            int count = in.readInt();
            List<FileStatus> files = new ArrayList<>();
            for (int i = 0; i < count; i++) {
                files.add(new FileStatus(
                    in.readUTF(),      // path
                    in.readLong(),     // length
                    in.readBoolean(),  // isDirectory
                    in.readInt(),      // replication
                    in.readLong()      // blockSize
                ));
            }
            return files;
        }
    }
}